# Copyright (C) 2024 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

import argparse
import json
import os
import re
from comps import ChatQnAGateway, MicroService, ServiceOrchestrator, ServiceType
from langchain_core.prompts import PromptTemplate


class ChatTemplate:
    # 生成RAG prompt
    @staticmethod
    def generate_rag_prompt(question, documents):
        context_str = "\n".join(documents)   # 将文档列表合并为上下文字符串
        if context_str and len(re.findall("[\u4E00-\u9FFF]", context_str)) / len(context_str) >= 0.3:
            # 判断上下文是否主要为中文内容，使用中文模板
            template = """
### 判断用户要求生成什么类型的图，根据用户要求生成对应的图的mermaid格式的UML代码。
如果用户要求生成的是ER图，用户提供一个实体关系（ER）描述文本，模型根据该文本生成符合 Mermaid 格式的 ER 图代码。
如果输入的描述是明确的、清晰的，模型应直接生成 ER 图代码,
如果描述较为抽象或信息不完全，模型需要通过推测、补全或简单的逻辑推理，给出可能的 ER 图代码。
ER 描述通常包括实体、属性以及它们之间的关系（例如，“实体A与实体B是一对多关系”）。
如果没有明确的描述或信息不完整，模型需要尽量推测合理的关系和属性，并输出相应的 Mermaid 代码。
例如，用户输入为：
“学生和课程之间是多对多关系，学生有姓名和学号属性，课程有课程名称和学分属性”，
输出的结果为：
erDiagram
STUDENT {
    string name
    string student_id
}
COURSE {
    string course_name
    int credits
}
STUDENT ||--o| COURSE : enrolls

如果用户要求生成的是时序图，用户提供一个时序图的描述文本，模型根据该文本生成符合 Mermaid 格式的时序图 UML 代码。
如果输入的描述是明确的、清晰的，模型应直接生成时序图代码；
如果描述较为抽象或信息不完全，模型需要通过推测、补全或简单的逻辑推理，给出可能的时序图代码。
时序图描述通常包括参与者（对象或系统）、消息的传递顺序以及交互的逻辑（例如，“用户提交登录请求，系统验证密码，返回结果”）。
如果描述缺少某些信息，模型将尽量推测合理的交互过程，并输出相应的 Mermaid 代码。
例如，用户输入为：“用户提交登录请求，系统验证密码，数据库返回验证结果，系统返回登录结果”，输出的结果为：
sequenceDiagram
participant User
participant System
participant Database
User->>System: Submit login request
System->>Database: Validate credentials
Database-->>System: Return validation result
System-->>User: Return login result

如果用户要求生成的是流程图，用户提供一个流程描述文本，模型根据该文本生成符合 Mermaid 格式的流程图代码。
如果输入的描述是明确的、清晰的，模型应直接生成流程图代码；
如果描述较为抽象或信息不完全，模型需要通过推测、补全或简单的逻辑推理，给出可能的流程图代码。
流程描述通常包括事件或节点之间的关系（例如，“从A到B”）。如果没有明确的描述顺序或节点，模型需要尽量推测合理的关系，并输出相应的 Mermaid 代码。
    比如用户输入为“ 生成流程图，描述从A到B，再从B到C”，输出的结果为
    flowchart TD
        A --> B
        B --> C
    如果输入的信息比较抽象或不完全明确，我们将进行检索并输出可能的 Mermaid 格式的UML代码。
    
如果用户要求生成的是甘特图。如果用户提供的是明确的任务描述、时间范围及任务之间的关系，模型应直接生成符合 Mermaid 格式的甘特图代码。
如果描述不完整或抽象，模型将通过推测和补全生成一个合理的甘特图代码。
比如用户输入："项目包含两个任务：任务A从1月1日到1月10日，任务B从1月5日到1月15日，任务B依赖于任务A完成后开始。"
输出的结果为
gantt
    title 项目甘特图
    dateFormat  YYYY-MM-DD
    section 项目任务
    任务A :a1, 2024-01-01, 10d
    任务B :after a1  , 2024-01-05, 10d
    


当用户要求生成类图时，模型会根据用户提供的描述生成符合 Mermaid 格式的类图代码。如果用户提供了清晰的类描述，
包括类名、类的属性和方法，以及类之间的关系（例如继承、关联等），模型会直接生成相应的类图代码。然而，如果用户的描述较为抽象或信息不完全，模型会通过推测、补全或简单的逻辑推理，尽可能补充缺失的部分，并生成合理的类图代码。
具体来说，类的描述应包括类的名称、属性和方法，属性应明确列出数据类型，方法应说明方法名及参数类型（如果有）。
类之间的关系可以是继承、关联、聚合或组合等。如果用户描述了类之间的继承关系（如“类A继承类B”），
模型将通过箭头表示继承关系，例如 `ClassA <|-- ClassB`。如果没有明确的类间关系，模型会根据推测生成合理的关联关系（如 `ClassA -- ClassB`）或其他形式的联系。
在描述不完整的情况下，模型会尽力推测类之间可能的关系并进行合理补充。
例如，如果用户只提供了部分类的属性或方法，模型会假设可能缺失的内容，并以合适的方式补充这些信息。最终，模型会输出符合 Mermaid 语法的类图代码，以便用户查看系统的类结构。
总之，明确的类描述会直接生成类图代码，抽象或不完全的描述则会通过推测和补全生成合理的类图代码。
比如用户输入为"系统有一个 Vehicle 类和一个 Car 类，Car 类继承自 Vehicle 类，Car 类有 brand 属性和 speed 属性。"
输出结果为：
classDiagram
    class Vehicle {
        +start()
        +stop()
    }

    class Car {
        +string brand
        +int speed
        +drive()
    }

    Vehicle <|-- Car : inherits

当用户要求生成饼图时，模型会根据用户提供的数据类别和相应的数值生成符合 
Mermaid 格式的饼图代码。如果用户提供了明确的数据描述，包括每个类别的名称和对应的百分比或数值，
模型会直接生成相应的饼图代码。用户还可以指定饼图的标题、标签、颜色等可视化配置。如果描述较为抽象或信息不完全，模型会推测或补充合理的数据，并生成合适的饼图代码。
具体来说，用户需要提供每个类别的名称以及该类别对应的数值或百分比。
如果用户描述了完整的数据（例如，“苹果占30%，香蕉占50%，橙子占20%”），
模型会直接生成相应的饼图代码，并可选地添加图表的标题。若用户描述不够具体或缺乏完整数据，
模型将推测合理的内容并补充相应的部分，生成一个合理的饼图。例如，若用户只提供了部分数据或类别，
模型会根据常见的情况做出合理假设，并继续生成合适的代码。
比如用户输入为"饼图显示了三种销售产品的占比，产品A占40%，产品B占30%，产品C占30%。"时，输出为：
pie
    title Sales Distribution
    "Product A" : 40
    "Product B" : 30
    "Product C" : 30


当用户要求生成Git图时，模型会根据用户提供的分支操作和提交记录生成符合Mermaid格式的Git图代码。Git图通常用来表示分支的创建、提交记录、合并操作以及分支之间的关系。用户需要提供清晰的描述，包括分支名称、提交操作的顺序、分支的创建、合并等操作。例如，用户可以描述“从master分支创建feature分支，feature分支上做了两次提交，然后将feature分支合并回master分支”，模型将根据这些描述生成Git图。如果用户没有明确描述某些操作，模型将推测合理的Git操作顺序并补充相应的信息，生成合适的Git图代码。
在具体操作上，用户需要描述分支的创建、提交、切换和合并的顺序。每次提交操作通过 commit 表示，创建新分支通过 branch <branch_name> 完成，切换分支则用 checkout <branch_name> 表示，合并分支用 merge <branch_name> 实现。如果描述比较抽象或信息不完整，模型会根据常见的Git操作推测合理的操作顺序，并生成Git图。用户也可以选择自定义Git图的标题、节点标签、颜色等可视化配置。
例如，若用户描述了“在main分支上做了两次提交，然后从main分支创建了dev分支，dev分支上做了两次提交，并最终将dev分支合并回main分支”，模型将生成对应的Git图代码，显示从main分支到dev分支的切换、提交记录以及最终的合并操作。总结来说，明确的Git操作描述将直接生成Git图代码，而较为抽象或不完整的描述将通过推测和补全生成合理的Git图。
比如用户输入为："从master分支创建feature分支，在feature分支上做了两次提交，然后将feature分支合并回master分支。"
输出为；
gitGraph
    commit
    branch feature
    commit
    commit
    checkout master
    merge feature

当用户要求生成用户旅程图时，模型会根据用户提供的描述生成符合Mermaid格式的用户旅程图代码。用户旅程图通常用来展示用户在与产品或服务的交互过程中经历的不同阶段、活动、情感变化以及这些阶段之间的关系。用户需要提供清晰的描述，涵盖用户旅程的不同阶段、每个阶段的具体活动或事件、用户在各阶段的情感变化以及这些阶段之间的顺序和转换关系。如果用户描述了明确的用户行为和情感变化，模型将直接生成相应的用户旅程图代码。如果描述较为抽象或不完整，模型将根据常见的用户旅程模型推测合理的阶段、活动和情感变化，并生成合适的用户旅程图代码。
具体来说，用户需要提供以下信息：首先是用户旅程的阶段，即用户在与产品或服务交互时所经历的主要阶段（例如“了解”、“考虑”、“购买”、“使用”等）；其次是每个阶段的活动或事件，描述用户在每个阶段进行的主要操作（例如浏览网站、添加购物车、支付等）；然后是每个阶段的用户情感变化，可以通过情感描述（如“高兴”、“沮丧”、“困惑”等）或情感标注（如“正面”、“负面”、“中立”）表示；最后是阶段之间的关系和转移，即描述用户如何从一个阶段转移到另一个阶段（如从“了解”到“考虑”）。如果用户提供了清晰的用户旅程描述，模型会直接生成符合Mermaid语法的用户旅程图代码。如果描述较为抽象或信息不完全，模型将推测合理的用户行为模式，并生成适当的旅程图。
例如，如果用户描述了“用户从了解产品开始，感到好奇；然后进入考虑阶段，感到困惑；接着进入购买阶段，感到高兴；最后进入使用阶段，感到满意”，模型会生成一张展示这些阶段和情感变化的用户旅程图。如果用户的描述不够具体，模型会通过推测和补全生成一个合理的用户旅程图，展示用户从开始浏览到最终使用的情感和行为路径。总的来说，明确的用户旅程描述将直接生成用户旅程图代码，而较为抽象或不完整的描述将通过推测和补充生成合适的旅程图。
比如用户输出为：“用户从了解产品开始，感到好奇；然后进入考虑阶段，感到困惑；接着进入购买阶段，感到高兴；最后进入使用阶段，感到满意。”
输出为：
journey
    title 用户旅程图
    section 了解产品
      用户感到好奇: 5: 产品介绍
    section 考虑购买
      用户感到困惑: 3: 比较选项
    section 购买产品
      用户感到高兴: 4: 完成购买
    section 使用产品
      用户感到满意: 5: 开始使用

当用户要求生成象限图时，模型会根据用户提供的数据或分类标准生成符合Mermaid格式的象限图代码。象限图通常用于展示两个变量或维度之间的关系，并将数据分布在四个象限中。用户需要提供清晰的描述，涵盖数据的两个维度，以及如何将数据点分类到不同的象限。如果用户提供了明确的维度和数据点，模型将直接生成相应的象限图代码。如果描述较为抽象或不完整，模型会根据常见的象限图数据模式推测合理的分类，并生成合适的象限图代码。
具体来说，用户需要提供以下信息：首先是用于分类数据的**两个维度**，例如X轴和Y轴的内容，可能是“价格”和“质量”或“市场份额”和“增长率”等；接下来是**象限区分标准**，即每个象限的分类标准，通常基于数据的中位数或用户定义的阈值，将数据分为四个区域，如“高-高”、“低-高”、“高-低”、“低-低”；然后是**数据点或类别**，用户需要提供数据点的名称以及它们对应的坐标，或者描述数据点在某个象限中的位置；最后，用户可以为象限图提供**标题和标签**，帮助解释数据。
如果用户提供了完整的维度和数据点，模型将直接生成符合Mermaid语法的象限图代码；如果描述不够具体或缺少完整数据，模型将推测合理的分类标准并生成合适的象限图。例如，用户可以描述“产品A（价格：高，质量：高）”，“产品B（价格：低，质量：高）”，“产品C（价格：低，质量：低）”，“产品D（价格：高，质量：低）”，模型会根据这些描述生成一个象限图，展示不同产品的分布。如果用户的描述较为抽象，模型会通过推测合理的维度和标准，生成合适的象限图。总之，明确的描述将直接生成象限图代码，而较为抽象或不完整的描述则会通过推测和补充生成适当的图表。
比如用户输入为：“我们有两项指标：X轴是‘产品价格’，Y轴是‘产品质量’，数据点包括‘产品A’（价格：高，质量：高），‘产品B’（价格：低，质量：高），‘产品C’（价格：低，质量：低），‘产品D’（价格：高，质量：低）。”
输出为；
quadrantChart
    title 产品象限图
    x-axis 价格
    y-axis 质量
    "产品A": [1, 1]
    "产品B": [-1, 1]
    "产品C": [-1, -1]
    "产品D": [1, -1]

### 搜索结果：{context}
### 问题：{question}
### 回答：
"""
        else:
            # 使用英文模板
            template = """
### Determine what type of diagram the user requests to generate, and generate the corresponding UML code in Mermaid format for the diagram according to the user's request.

If the user requests to generate an ER diagram, the user provides an entity-relationship (ER) description text, and the model generates the ER diagram code in Mermaid format based on this text.
If the input description is clear and explicit, the model should directly generate the ER diagram code.
If the description is rather abstract or the information is incomplete, the model needs to infer, complete, or use simple logical reasoning to give possible ER diagram code.
The ER description usually includes entities, attributes, and the relationships between them (for example, "Entity A has a one-to-many relationship with Entity B").
If there is no clear description or the information is incomplete, the model needs to try its best to infer reasonable relationships and attributes and output the corresponding Mermaid code.
For example, if the user's input is:
"There is a many-to-many relationship between students and courses. Students have the attributes of name and student ID, and courses have the attributes of course name and credits.",
the output result is:
erDiagram
    STUDENT {
        string name
        string student_id
    }
    COURSE {
        string course_name
        int credits
    }
    STUDENT ||--o| COURSE : enrolls

If the user requests to generate a sequence diagram, the user provides a description text of the sequence diagram, and the model generates the sequence diagram UML code in Mermaid format based on this text.
If the input description is clear and explicit, the model should directly generate the sequence diagram code.
If the description is rather abstract or the information is incomplete, the model needs to infer, complete, or use simple logical reasoning to give possible sequence diagram code.
The sequence diagram description usually includes participants (objects or systems), the order of message transmission, and the logic of interaction (for example, "The user submits a login request, the system verifies the password, and returns the result").
If the description lacks certain information, the model will try its best to infer a reasonable interaction process and output the corresponding Mermaid code.
For example, if the user's input is: "The user submits a login request, the system verifies the password, the database returns the verification result, and the system returns the login result.", the output result is:
sequenceDiagram
    participant User
    participant System
    participant Database
    User->>System: Submit login request
    System->>Database: Validate credentials
    Database-->>System: Return validation result
    System-->>User: Return login result

If the user requests to generate a flowchart, the user provides a process description text, and the model generates the flowchart code in Mermaid format based on this text.
If the input description is clear and explicit, the model should directly generate the flowchart code.
If the description is rather abstract or the information is incomplete, the model needs to infer, complete, or use simple logical reasoning to give possible flowchart code.
The process description usually includes the relationships between events or nodes (for example, "From A to B"). If there is no clear description order or nodes, the model needs to try its best to infer reasonable relationships and output the corresponding Mermaid code.
For example, if the user's input is: "Generate a flowchart to describe going from A to B and then from B to C.", the output result is:
flowchart TD
    A --> B
    B --> C
    
If the user requests to generate a Gantt chart, and the user provides a clear task description, time range, and relationships between tasks, the model should directly generate the corresponding Mermaid format Gantt chart code. If the description is incomplete or abstract, the model should infer and complete the information to generate a reasonable Gantt chart code.
For example, if the user inputs:
"The project contains two tasks: Task A starts on January 1st and ends on January 10th, and Task B starts on January 5th and ends on January 15th. Task B depends on Task A to finish before it can start."
The output would be:
gantt
    title Project Gantt Chart
    dateFormat  YYYY-MM-DD
    section Project Tasks
    Task A :a1, 2024-01-01, 10d
    Task B :after a1  , 2024-01-05, 10d
    
When the user requests to generate a class diagram, the model will generate the corresponding Mermaid format class diagram code based on the description provided by the user. If the user provides a clear description of the classes, including class names, attributes, methods, and relationships between classes (such as inheritance, association, etc.), the model will directly generate the appropriate class diagram code. However, if the description is more abstract or incomplete, the model will use inference, supplementation, or simple logical reasoning to fill in the missing details and generate a reasonable class diagram code.
Specifically, the description of the class should include the class name, attributes, and methods. The attributes should be clearly listed with data types, and the methods should specify method names and parameter types (if any). Relationships between classes can include inheritance, association, aggregation, or composition. If the user describes an inheritance relationship between classes (e.g., "Class A inherits from Class B"), the model will represent this inheritance with an arrow, such as ClassA <|-- ClassB. If no explicit relationship between classes is described, the model will infer a reasonable association (e.g., ClassA -- ClassB) or other forms of connections.
In cases of incomplete descriptions, the model will try to infer possible relationships between classes and make reasonable additions. For example, if the user only provides some attributes or methods of a class, the model will assume what might be missing and supplement the information accordingly. Ultimately, the model will output class diagram code in Mermaid syntax, allowing the user to visualize the structure of the system's classes.
In summary, a clear class description will result in direct class diagram code, while abstract or incomplete descriptions will be supplemented and inferred to generate reasonable class diagram code.
For example, if the user inputs:
"The system has a Vehicle class and a Car class. The Car class inherits from the Vehicle class, and the Car class has brand and speed attributes."
The output would be:
classDiagram
    class Vehicle {
        +start()
        +stop()
    }

    class Car {
        +string brand
        +int speed
        +drive()
    }

    Vehicle <|-- Car : inherits

When the user requests to generate a pie chart, the model will generate the corresponding Mermaid format pie chart code based on the data categories and their corresponding values provided by the user. If the user provides a clear data description, including the name of each category and its corresponding percentage or value, the model will directly generate the appropriate pie chart code. Users can also specify additional visual configurations such as the title, labels, and colors of the pie chart. If the description is more abstract or incomplete, the model will infer or supplement reasonable data and generate a suitable pie chart code.
Specifically, the user needs to provide the name of each category and the corresponding value or percentage. If the user describes complete data (e.g., "Apple 30%, Banana 50%, Orange 20%"), the model will directly generate the corresponding pie chart code, optionally adding a title for the chart. If the description is not specific enough or lacks complete data, the model will infer reasonable content and fill in the missing parts, generating a suitable pie chart. For example, if the user only provides partial data or categories, the model will make reasonable assumptions based on common patterns and continue to generate appropriate code.
For example, if the user inputs: "The pie chart shows the market share of three products: Product A 40%, Product B 30%, Product C 30%", the output will be:
pie
    title Sales Distribution
    "Product A" : 40
    "Product B" : 30
    "Product C" : 30


When the user requests to generate a Git graph, the model will generate the corresponding Mermaid format Git graph code based on the branch operations and commit records provided by the user. A Git graph is typically used to represent branch creation, commit records, merge operations, and relationships between branches. The user needs to provide a clear description, including branch names, the order of commit operations, branch creation, and merging operations. For example, if the user describes "creating a `feature` branch from `master`, making two commits on the `feature` branch, and then merging `feature` back into `master`," the model will generate a Git graph based on this description. If the user does not explicitly describe certain operations, the model will infer a reasonable Git operation sequence and fill in the missing information to generate appropriate Git graph code.
Specifically, the user needs to describe the sequence of branch creation, commits, checkouts, and merges. Each commit operation is represented by `commit`, creating a new branch is done with `branch <branch_name>`, switching branches is represented by `checkout <branch_name>`, and merging a branch is done with `merge <branch_name>`. If the description is more abstract or incomplete, the model will infer a reasonable sequence of Git operations and generate a suitable Git graph. Users can also customize the title, node labels, colors, and other visual configurations of the Git graph.
For example, if the user describes "making two commits on the `main` branch, then creating a `dev` branch from `main`, making two commits on `dev`, and finally merging `dev` back into `main`," the model will generate the corresponding Git graph code, showing the switch from `main` to `dev`, the commits, and the final merge operation. In summary, clear Git operation descriptions will directly generate Git graph code, while more abstract or incomplete descriptions will be inferred and supplemented to generate a reasonable Git graph.
For example, if the user inputs: "Create a `feature` branch from `master`, make two commits on the `feature` branch, and then merge `feature` back into `master`," the output will be:
If the input information is rather abstract or not completely clear, we will conduct a search and output the possible UML code in Mermaid format. 
gitGraph
    commit
    branch feature
    commit
    commit
    checkout master
    merge feature

When the user requests to generate a user journey map, the model will generate the corresponding Mermaid format user journey map code based on the user's description. A user journey map is typically used to display the different stages, activities, emotional changes, and relationships between stages that a user experiences while interacting with a product or service. The user needs to provide a clear description, covering the different stages of the user journey, the specific activities or events in each stage, the emotional changes of the user at each stage, and the order and transitions between these stages. If the user describes clear user behaviors and emotional changes, the model will directly generate the corresponding user journey map code. If the description is more abstract or incomplete, the model will infer reasonable stages, activities, and emotional changes based on common user journey models and generate an appropriate user journey map code.
Specifically, the user needs to provide the following information: First, the stages of the user journey, which represent the key phases the user goes through when interacting with the product or service (e.g., "Awareness," "Consideration," "Purchase," "Usage," etc.); Second, the activities or events in each stage, describing the main actions the user performs at each stage (e.g., browsing the website, adding to cart, making a payment, etc.); Third, the emotional changes of the user in each stage, which can be expressed through emotional descriptions (e.g., "happy," "frustrated," "confused," etc.) or emotional labels (e.g., "positive," "negative," "neutral"); Lastly, the relationships and transitions between stages, describing how the user moves from one stage to another (e.g., from "Awareness" to "Consideration"). If the user provides a clear description of the user journey, the model will directly generate the user journey map code in Mermaid syntax. If the description is more abstract or incomplete, the model will infer a reasonable user behavior pattern and generate an appropriate journey map.
For example, if the user describes: "The user starts by becoming aware of the product, feeling curious; then moves to the consideration stage, feeling confused; next, moves to the purchase stage, feeling happy; and finally, moves to the usage stage, feeling satisfied," the model will generate a user journey map that shows these stages and emotional changes. If the user's description is not specific enough, the model will infer and fill in reasonable details to create a user journey map that shows the emotional and behavioral path from the start of browsing to the final usage. In summary, a clear description of the user journey will directly generate the user journey map code, while more abstract or incomplete descriptions will be inferred and supplemented to generate an appropriate journey map.
For example, if the user inputs: "The user starts by becoming aware of the product, feeling curious; then moves to the consideration stage, feeling confused; next, moves to the purchase stage, feeling happy; and finally, moves to the usage stage, feeling satisfied," the output would be:
journey
    title User Journey Map
    section Awareness of Product
      User feels curious: 5: Product Introduction
    section Consideration
      User feels confused: 3: Compare Options
    section Purchase
      User feels happy: 4: Complete Purchase
    section Usage
      User feels satisfied: 5: Start Using

When a user requests to generate a quadrant chart, the model will generate the corresponding Mermaid-format quadrant chart code based on the data or classification criteria provided by the user. A quadrant chart is typically used to show the relationship between two variables or dimensions and to distribute data across four quadrants. The user needs to provide a clear description, covering the two dimensions of the data and how the data points are classified into different quadrants. If the user provides explicit dimensions and data points, the model will directly generate the appropriate quadrant chart code. If the description is abstract or incomplete, the model will infer a reasonable classification based on common quadrant chart patterns and generate a suitable quadrant chart code.
Specifically, the user needs to provide the following information: first, the two dimensions used to classify the data, such as "price" and "quality" or "market share" and "growth rate"; next, the quadrant classification criteria, which is usually based on the median of the data or user-defined thresholds, dividing the data into four regions like "high-high," "low-high," "high-low," and "low-low"; then, the data points or categories, where the user provides the name of each data point and its corresponding coordinates or describes its position in a particular quadrant; finally, the user can provide a title and labels for the quadrant chart to help explain the data.
If the user provides complete dimensions and data points, the model will directly generate the corresponding Mermaid-format quadrant chart code. If the description is unclear or missing key data, the model will infer reasonable classification criteria and generate an appropriate quadrant chart. For example, the user might describe "Product A (Price: high, Quality: high)," "Product B (Price: low, Quality: high)," "Product C (Price: low, Quality: low)," "Product D (Price: high, Quality: low)," and the model will generate a quadrant chart showing the distribution of these products. If the user’s description is abstract, the model will infer reasonable dimensions and standards and generate an appropriate chart. In conclusion, explicit descriptions will directly generate quadrant chart code, while more abstract or incomplete descriptions will be inferred and supplemented to generate an appropriate chart.
For example, if the user inputs: "We have two metrics: the X-axis is 'Product Price' and the Y-axis is 'Product Quality,' with data points including 'Product A' (Price: high, Quality: high), 'Product B' (Price: low, Quality: high), 'Product C' (Price: low, Quality: low), 'Product D' (Price: high, Quality: low)," the output would be:
quadrantChart
    title Product Quadrant Chart
    x-axis Price
    y-axis Quality
    "Product A": [1, 1]
    "Product B": [-1, 1]
    "Product C": [-1, -1]
    "Product D": [1, -1]

### Search results: {context} \n
### Question: {question} \n
### Answer:
"""
        # 格式化模板
        return template.format(context=context_str, question=question)

# 定义服务的默认主机和端口号
MEGA_SERVICE_HOST_IP = os.getenv("MEGA_SERVICE_HOST_IP", "0.0.0.0")
MEGA_SERVICE_PORT = int(os.getenv("MEGA_SERVICE_PORT", 8888))
GUARDRAIL_SERVICE_HOST_IP = os.getenv("GUARDRAIL_SERVICE_HOST_IP", "0.0.0.0")
GUARDRAIL_SERVICE_PORT = int(os.getenv("GUARDRAIL_SERVICE_PORT", 80))
EMBEDDING_SERVER_HOST_IP = os.getenv("EMBEDDING_SERVER_HOST_IP", "0.0.0.0")
EMBEDDING_SERVER_PORT = int(os.getenv("EMBEDDING_SERVER_PORT", 80))
RETRIEVER_SERVICE_HOST_IP = os.getenv("RETRIEVER_SERVICE_HOST_IP", "0.0.0.0")
RETRIEVER_SERVICE_PORT = int(os.getenv("RETRIEVER_SERVICE_PORT", 7000))
RERANK_SERVER_HOST_IP = os.getenv("RERANK_SERVER_HOST_IP", "0.0.0.0")
RERANK_SERVER_PORT = int(os.getenv("RERANK_SERVER_PORT", 80))
LLM_SERVER_HOST_IP = os.getenv("LLM_SERVER_HOST_IP", "0.0.0.0")
LLM_SERVER_PORT = int(os.getenv("LLM_SERVER_PORT", 80))
LLM_MODEL = os.getenv("LLM_MODEL", "Intel/neural-chat-7b-v3-3")

# 对输入进行对齐和调整
# 根据当前节点的服务类型，调整输入数据的格式
# 包括EMBEDDING、RETRIEVER、LLM等服务的输入格式调整

def align_inputs(self, inputs, cur_node, runtime_graph, llm_parameters_dict, **kwargs):
    if self.services[cur_node].service_type == ServiceType.EMBEDDING:
        inputs["inputs"] = inputs["text"]  # 将输入文本字段转换为"inputs"
        del inputs["text"]  # 删除原始"text"字段
    elif self.services[cur_node].service_type == ServiceType.RETRIEVER:
        # 如果是RETRIEVER服务，更新检索器参数
        retriever_parameters = kwargs.get("retriever_parameters", None)
        if retriever_parameters:
            inputs.update(retriever_parameters.dict())
    elif self.services[cur_node].service_type == ServiceType.LLM:
        # 如果是LLM服务，构造特定的输入格式
        next_inputs = {}
        next_inputs["model"] = LLM_MODEL
        next_inputs["messages"] = [{"role": "user", "content": inputs["inputs"]}]
        next_inputs["max_tokens"] = llm_parameters_dict["max_tokens"]
        next_inputs["top_p"] = llm_parameters_dict["top_p"]
        next_inputs["stream"] = inputs["streaming"]
        next_inputs["frequency_penalty"] = inputs["frequency_penalty"]
        next_inputs["temperature"] = inputs["temperature"]
        inputs = next_inputs  # 更新输入格式
    return inputs

# 对输出进行对齐和调整
# 根据服务类型处理输出数据格式，包括嵌入、检索、重排序和LLM服务

def align_outputs(self, data, cur_node, inputs, runtime_graph, llm_parameters_dict, **kwargs):
    next_data = {}
    if self.services[cur_node].service_type == ServiceType.EMBEDDING:
        assert isinstance(data, list)  # 确保嵌入服务的输出为列表
        next_data = {"text": inputs["inputs"], "embedding": data[0]}  # 返回文本和嵌入结果
    elif self.services[cur_node].service_type == ServiceType.RETRIEVER:
        docs = [doc["text"] for doc in data["retrieved_docs"]]  # 提取检索到的文档
        with_rerank = runtime_graph.downstream(cur_node)[0].startswith("rerank")  # 检查是否需要重排序
        if with_rerank and docs:
            next_data["query"] = data["initial_query"]
            next_data["texts"] = docs
        else:
            if not docs and with_rerank:
                # 如果没有文档且需要重排序，调整运行图
                for ds in reversed(runtime_graph.downstream(cur_node)):
                    for nds in runtime_graph.downstream(ds):
                        runtime_graph.add_edge(cur_node, nds)
                    runtime_graph.delete_node_if_exists(ds)
            prompt = data["initial_query"]
            chat_template = llm_parameters_dict["chat_template"]
            if chat_template:
                # 根据模板生成提示
                prompt_template = PromptTemplate.from_template(chat_template)
                input_variables = prompt_template.input_variables
                if sorted(input_variables) == ["context", "question"]:
                    prompt = prompt_template.format(question=data["initial_query"], context="\n".join(docs))
                elif input_variables == ["question"]:
                    prompt = prompt_template.format(question=data["initial_query"])
                else:
                    print(f"{prompt_template} not used, we only support 2 input variables ['question', 'context']")
                    prompt = ChatTemplate.generate_rag_prompt(data["initial_query"], docs)
            else:
                prompt = ChatTemplate.generate_rag_prompt(data["initial_query"], docs)
            next_data["inputs"] = prompt
    elif self.services[cur_node].service_type == ServiceType.RERANK:
        # 处理RERANK服务的输出
        reranker_parameters = kwargs.get("reranker_parameters", None)
        top_n = reranker_parameters.top_n if reranker_parameters else 1
        docs = inputs["texts"]
        reranked_docs = []
        for best_response in data[:top_n]:
            reranked_docs.append(docs[best_response["index"]])
        prompt = inputs["query"]
        chat_template = llm_parameters_dict["chat_template"]
        if chat_template:
            prompt_template = PromptTemplate.from_template(chat_template)
            input_variables = prompt_template.input_variables
            if sorted(input_variables) == ["context", "question"]:
                prompt = prompt_template.format(question=prompt, context="\n".join(reranked_docs))
            elif input_variables == ["question"]:
                prompt = prompt_template.format(question=prompt)
            else:
                print(f"{prompt_template} not used, we only support 2 input variables ['question', 'context']")
                prompt = ChatTemplate.generate_rag_prompt(prompt, reranked_docs)
        else:
            prompt = ChatTemplate.generate_rag_prompt(prompt, reranked_docs)
        next_data["inputs"] = prompt
    elif self.services[cur_node].service_type == ServiceType.LLM and not llm_parameters_dict["streaming"]:
        next_data["text"] = data["choices"][0]["message"]["content"]  # 返回LLM生成的文本内容
    else:
        next_data = data  # 直接返回数据
    return next_data

# 对生成器数据流进行对齐
# 用于处理LLM服务的流式响应

def align_generator(self, gen, **kwargs):
    for line in gen:
        line = line.decode("utf-8")  # 解码流数据
        start = line.find("{")
        end = line.rfind("}") + 1
        json_str = line[start:end]
        try:
            # 解析JSON数据
            json_data = json.loads(json_str)
            if (
                json_data["choices"][0]["finish_reason"] != "eos_token"
                and "content" in json_data["choices"][0]["delta"]
            ):
                yield f"data: {repr(json_data['choices'][0]['delta']['content'].encode('utf-8'))}\n\n"
        except Exception as e:
            yield f"data: {repr(json_str.encode('utf-8'))}\n\n"
    yield "data: [DONE]\n\n"


class ChatService:
    # 使用各类微服务构建问答服务
    def __init__(self, host="0.0.0.0", port=8000):
        self.host = host  # 服务的主机地址
        self.port = port  # 服务的端口号
        # 将输入和输出对齐函数绑定到ServiceOrchestrator类
        ServiceOrchestrator.align_inputs = align_inputs
        ServiceOrchestrator.align_outputs = align_outputs
        ServiceOrchestrator.align_generator = align_generator
        # 初始化微服务编排器
        self.megaservice = ServiceOrchestrator()

    # 添加远程服务，包括embedding、retriever、rerank、llm
    def add_remote_service(self):

        # 定义embedding服务
        embedding = MicroService(
            name="embedding",
            host=EMBEDDING_SERVER_HOST_IP,
            port=EMBEDDING_SERVER_PORT,
            endpoint="/embed",
            use_remote_service=True,
            service_type=ServiceType.EMBEDDING,
        )

        # 定义retriever服务
        retriever = MicroService(
            name="retriever",
            host=RETRIEVER_SERVICE_HOST_IP,
            port=RETRIEVER_SERVICE_PORT,
            endpoint="/v1/retrieval",
            use_remote_service=True,
            service_type=ServiceType.RETRIEVER,
        )

        # 定义rerank服务
        rerank = MicroService(
            name="rerank",
            host=RERANK_SERVER_HOST_IP,
            port=RERANK_SERVER_PORT,
            endpoint="/rerank",
            use_remote_service=True,
            service_type=ServiceType.RERANK,
        )

        # 定义llm服务
        llm = MicroService(
            name="llm",
            host=LLM_SERVER_HOST_IP,
            port=LLM_SERVER_PORT,
            endpoint="/v1/chat/completions",
            use_remote_service=True,
            service_type=ServiceType.LLM,
        )

        # 添加微服务到编排器
        self.megaservice.add(embedding).add(retriever).add(rerank).add(llm)
        # 配置服务之间的调用流程
        self.megaservice.flow_to(embedding, retriever)
        self.megaservice.flow_to(retriever, rerank)
        self.megaservice.flow_to(rerank, llm)
        # 初始化网关服务
        self.gateway = ChatQnAGateway(megaservice=self.megaservice, host="0.0.0.0", port=self.port)

    # 添加远程服务（不包括rerank）
    def add_remote_service_without_rerank(self):

        # 定义embedding服务
        embedding = MicroService(
            name="embedding",
            host=EMBEDDING_SERVER_HOST_IP,
            port=EMBEDDING_SERVER_PORT,
            endpoint="/embed",
            use_remote_service=True,
            service_type=ServiceType.EMBEDDING,
        )

        # 定义retriever服务
        retriever = MicroService(
            name="retriever",
            host=RETRIEVER_SERVICE_HOST_IP,
            port=RETRIEVER_SERVICE_PORT,
            endpoint="/v1/retrieval",
            use_remote_service=True,
            service_type=ServiceType.RETRIEVER,
        )

        # 定义llm服务
        llm = MicroService(
            name="llm",
            host=LLM_SERVER_HOST_IP,
            port=LLM_SERVER_PORT,
            endpoint="/v1/chat/completions",
            use_remote_service=True,
            service_type=ServiceType.LLM,
        )

        # 添加微服务到编排器
        self.megaservice.add(embedding).add(retriever).add(llm)
        # 配置服务之间的调用流程
        self.megaservice.flow_to(embedding, retriever)
        self.megaservice.flow_to(retriever, llm)
        # 初始化网关服务
        self.gateway = ChatQnAGateway(megaservice=self.megaservice, host="0.0.0.0", port=self.port)

    # 添加远程服务并包含guardrail服务
    def add_remote_service_with_guardrails(self):
        # 定义guardrail服务
        guardrail_in = MicroService(
            name="guardrail_in",
            host=GUARDRAIL_SERVICE_HOST_IP,
            port=GUARDRAIL_SERVICE_PORT,
            endpoint="/v1/guardrails",
            use_remote_service=True,
            service_type=ServiceType.GUARDRAIL,
        )
        # 定义embedding服务
        embedding = MicroService(
            name="embedding",
            host=EMBEDDING_SERVER_HOST_IP,
            port=EMBEDDING_SERVER_PORT,
            endpoint="/embed",
            use_remote_service=True,
            service_type=ServiceType.EMBEDDING,
        )
        # 定义retriever服务
        retriever = MicroService(
            name="retriever",
            host=RETRIEVER_SERVICE_HOST_IP,
            port=RETRIEVER_SERVICE_PORT,
            endpoint="/v1/retrieval",
            use_remote_service=True,
            service_type=ServiceType.RETRIEVER,
        )
        # 定义rerank服务
        rerank = MicroService(
            name="rerank",
            host=RERANK_SERVER_HOST_IP,
            port=RERANK_SERVER_PORT,
            endpoint="/rerank",
            use_remote_service=True,
            service_type=ServiceType.RERANK,
        )
        # 定义llm服务
        llm = MicroService(
            name="llm",
            host=LLM_SERVER_HOST_IP,
            port=LLM_SERVER_PORT,
            endpoint="/v1/chat/completions",
            use_remote_service=True,
            service_type=ServiceType.LLM,
        )

        # 添加微服务到编排器
        self.megaservice.add(guardrail_in).add(embedding).add(retriever).add(rerank).add(llm)
        # 配置服务之间的调用流程
        self.megaservice.flow_to(guardrail_in, embedding)
        self.megaservice.flow_to(embedding, retriever)
        self.megaservice.flow_to(retriever, rerank)
        self.megaservice.flow_to(rerank, llm)
        # 初始化网关服务
        self.gateway = ChatQnAGateway(megaservice=self.megaservice, host="0.0.0.0", port=self.port)

# 主程序入口
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # 添加参数选项：是否不包含rerank服务
    parser.add_argument("--without-rerank", action="store_true")
    # 添加参数选项：是否包含guardrails服务
    parser.add_argument("--with-guardrails", action="store_true")

    # 解析命令行参数
    args = parser.parse_args()

    # 初始化ChatQnA服务
    chat = ChatService(host=MEGA_SERVICE_HOST_IP, port=MEGA_SERVICE_PORT)
    if args.without_rerank:
        chat.add_remote_service_without_rerank()  # 不包含rerank服务
    elif args.with_guardrails:
        chat.add_remote_service_with_guardrails()  # 包含guardrails服务
    else:
        chat.add_remote_service()  # 默认服务配置